#include "../include/PageProcessor.h"

using std::cout;
using std::string;
using std::regex;
using std::endl;
using std::vector;
using std::set;
using std::ifstream;
using std::cerr;
using std::ofstream;
using std::unordered_set;

using namespace tinyxml2;
using namespace simhash;

PageProcessor::PageProcessor()
{
    ifstream ifs{"../corpus/stopwords/cn_stopwords.txt"};
    string line;
    while(getline(ifs, line))
    {
        istringstream iss{line};
        string word;
        while(iss >> word)
        {
            m_stopWords.insert(word);
        }
    }
    ifs.close(); 
}

void PageProcessor::process(const std::string& dir)
{
    vector<string>Corpus = DirectoryScanner::scan(dir);
    int i = 0;
    for(auto &e : Corpus)
    {
        extract_documents(e);
        i++;
    }
    deduplicate_documents();

    build_pages_and_offsets("../libfile/pageslib.dat", "../libfile/offsetslib.dat");

    build_inverted_index("../libfile/invertedIndexLib.dat");
}

void PageProcessor::extract_documents(const std::string& dir)
{
    XMLDocument doc;
    doc.LoadFile(dir.c_str());
    if(doc.ErrorID())
    {
        cerr << "LoadFile fail" << endl;
    }

    XMLElement *itemNode = doc.FirstChildElement("rss")
                             ->FirstChildElement("channel")
                             ->FirstChildElement("item");
   
    while(itemNode)
    {
        Document document;
        if(auto linkNode = itemNode->FirstChildElement("link"))
        {
            document.link = linkNode->GetText();
        }

        if(auto titleNode = itemNode->FirstChildElement("title"))
        {
            document.title = titleNode->GetText();
        }

        if(auto contentNode = itemNode->FirstChildElement("content:encoded"))
        {
            document.content = contentNode->GetText();
        }

        if(document.content.empty())
        {
            if(auto contentNode = itemNode->FirstChildElement("description"))
            {
                document.content = contentNode->GetText();
            }
            if(document.content.empty())
            {
                itemNode = itemNode->NextSiblingElement("item");
                continue;
            }
        }

        regex reg{"<[^>]+>"};
        document.content = std::regex_replace(document.content, reg, "");

        m_documents.push_back(document);
        itemNode = itemNode->NextSiblingElement("item");
    }
}

void PageProcessor::deduplicate_documents()
{
    vector<uint64_t> hashCodeLists;
    for(auto &e : m_documents)
    {
        size_t topN;
        if(e.content.size() < 30)
        {
            topN = 10;
        }else if(e.content.size() >= 30 && e.content.size() < 100)
        {
            topN = 30;
        }else if(e.content.size() >= 100 && e.content.size() < 300)
        {
            topN = 100;
        }else
        {
            topN = 300;
        }

        uint64_t hashCode;

        m_hasher.make(e.content, topN, hashCode);
        hashCodeLists.push_back(hashCode);
    }

    // 去重
    unordered_set<uint64_t> noSimilar;
    vector<Document> dedup;
    for(size_t i = 0; i < m_documents.size(); i++)
    {
        if(noSimilar.insert(hashCodeLists[i]).second)
        {
            dedup.push_back(m_documents[i]);
        }
    }
    m_documents.swap(dedup);
}

void PageProcessor::build_pages_and_offsets(const std::string& pages, const std::string& offsets)
{
    ofstream ofs1{pages};
    ofstream ofs2{offsets};
    size_t idx = 1; 
    off_t pageOffset = 0;
    for(auto &e : m_documents)
    {
        e.id = idx;
       
        auto start = ofs1.tellp();

        ofs1 << "<doc>" << endl
             << "  <id>" << e.id << "</id>" << endl
             << "  <link>" << e.link << "</link>" << endl
             << "  <title>" << e.title << "</title>" << endl
             << "  <content>" << e.content << "</content>" << endl
             << "</doc>" << endl;

        auto end = ofs1.tellp();
        size_t size = (size_t)(end - start);
        ofs2 << idx << " " << pageOffset << " " << size << endl;

        pageOffset += size;
        idx++;
    }

    ofs1.close();
    ofs2.close();
}

void PageProcessor::build_inverted_index(const std::string& filename)
{
    vector<string> words;
    map<int, int> wordSums; // 第i篇文档的总词数
    map<int, map<string, int>> wordFrequency; // 第i篇文档中的关键字，以及该关键字在该文档出现的次数
    
    for(const auto &doc : m_documents)
    {
        m_tokenizer.Cut(doc.content, words);
        wordSums[doc.id] = words.size();
        for(const auto &word : words)
        {
            auto it = utf8::iterator
            {
                word.begin(), 
                word.begin(), 
                word.end()
            };
            char32_t codePoint = *it;
            if(codePoint < 0x4E00 || codePoint > 0x9FFF)
            {
                continue;
            }
            if(m_stopWords.find(word) == m_stopWords.end())
            {
                m_invertedIndex[word][doc.id];
                ++wordFrequency[doc.id][word];
            }
        }     
    } 

    map<int, map<string, double>> forwardIndex; // 正排索引，保存第i篇文档的所有关键字及其权重
    
    // 计算TF-IDF权重
    for(auto &keyword : m_invertedIndex)
    {
        for(auto &[id, weight] : keyword.second)
        {
            double TF = static_cast<double>(wordFrequency[id][keyword.first]) 
                      / static_cast<double>(wordSums[id]);
            double IDF = std::log2(
                         (static_cast<double>(m_documents.size()) + 1.0)
                       / (static_cast<double>(keyword.second.size()) + 1.0)
            );
            weight = TF * IDF;
            forwardIndex[id][keyword.first] = weight;
        }
    }

    // 归一化处理
    for(auto &keyword : m_invertedIndex)
    {
        for(auto &[id, weight] : keyword.second)
        {
            double sum_of_squares;
            for(auto &ele : forwardIndex[id])
            {
                sum_of_squares += (ele.second * ele.second);
            }
            weight = weight / sqrt(sum_of_squares);
        }
    }   

    ofstream ofs{filename};
    for(const auto &keyWord : m_invertedIndex)
    {
        ofs << keyWord.first << " ";
        for(const auto &[id, weigth] : keyWord.second)
        {
            ofs << id << " " << weigth << " ";
        }
        ofs << endl;
    }
    ofs.close();
}
